#'
#' Process the mlogit formula and generate a list 
#' of three formula, representing the formula for 
#' individual specific variables that must have coefficients 
#' that vary by alternatives, alternative specific variables 
#' that have a common constant coeffient, and alternative 
#' specific variables that have alternative dependent coefficients. 
#'
#' @return a list of three formula
#' @note For the second and third components, no intercepts 
#' are allowed. 
findBars <- function(formula){
  ftext <- strsplit(deparse(formula), "~")
  ftext <- sapply(ftext, function(x) x[length(x)])
  if (length(ftext) > 1)
    ftext <- paste0(ftext, collapse = "")
  terms <- strsplit(ftext, "|", fixed = TRUE)
  terms <- gsub("\\s", "", terms[[1]])
  nt <- length(terms)
  forms <- lapply(1:nt, function(t){
    x <- terms[[t]]
    form <- NULL
    if (x != ""){
      form <- as.formula(paste("~", paste(x, collapse = "+")))
      #if (t > 1)
      #  form <- update(form, "~. + 0")
      environment(form) <- environment(formula)
    }    
    form
  })    
  if (nt < 3)
    forms <- c(forms, rep(list(NULL), 3 - nt))
  return(forms)
}

#' Remove the separating bars in an input formula
noBars <- function(formula){
  ftext <- strsplit(deparse(formula), "~")
  ftext <- sapply(ftext, function(x) x[length(x)])
  if (length(ftext) > 1)
    ftext <- paste0(ftext, collapse = "")
  terms <- strsplit(ftext, "|", fixed = TRUE)
  terms <- gsub("\\s", "", terms[[1]])
  terms <- terms[terms != ""]
  update(formula, paste(".~", paste(terms, collapse = "+")))
}

# find the external grouping structure
findGroup <- function(term) {
  if (is.numeric(term))
    return(NULL)
  if (is.name(term) || !is.language(term)) 
    return(term)
  if (term[[1]] == as.name("(")) 
    return(term[[2]])
  stopifnot(is.call(term))
  if (term[[1]] == as.name("|")) 
    return(term)
  if (length(term) == 2) 
    return(findGroup(term[[2]]))
  c(findGroup(term[[2]]), findGroup(term[[3]]))
}

# get the mapping from variable position to group position
getAssign <- function(formula){
  grp <- findGroup(formula)
  grp <- sapply(grp, function(x) strsplit(deparse(x), "\\+"))
  len <- cumsum(sapply(grp, length))
  assign <- seq_len(max(len))
  assign <- as.integer(cut(assign, c(0, len)))
  return(assign)
}

# check input data 
check.data <- function(data){
  if (!("mlogit.data" %in% class(data)))
    stop("Input data must be of class 'mlogit.data'.")
  idx <- index(data)
  if (nlevels(idx$chid) * nlevels(idx$alt) != nrow(data))
    stop("Input data is not balanced.")
}    

#'
#' Get model frame and design matrix.
#'
#' @param mc matched call 
#' @param formula the input model formula
#' @param use.contrast whether the zero-based contrast is used. If true, 
#'  the first base level is set to zero. Otherwise, all dummys are used. 
#'  Other contrast methods are not supported for now. 
#' 
#' @return a list of the following elements:
#' \item{mf} model frame containing all relevant covariates
#' \item{X} the sparse design matrix in column oriented form. 
#' It has an attribute "grp" which is used to determine the grouping
#' structure of the coefficients corresponding to the first formula
#' element. This is used to center the estimates after each 
#' coordinate descent step. It's 0 based.
#' \item{y} the (integer) response. It has an attribute "choice"
#' that indicates which alternative is selected, and a second 
#' attribute "nzero" that indicates which row of y is nonzero (0 based). 
#' \item{alt} the factor vector of alternative indicator
#' \item{id} the factor vector of individual indicator
#' \item{formula} the input formula
#' \item{dist.vars} the covariates that are used to compute the 
#' distances among individuals. A vector of characters. 

getFrame <- function(mc, formula, use.contrast = TRUE, 
                     drop.unused.levels = TRUE, 
                     xlev = NULL) {
  mff <- mc
  bars <- findBars(formula)
  frame.form <- noBars(formula)

  # 1. get model frame 
  # we do not allow weights or offset now
  m <- match(c("data", "subset", "weights", "na.action"), names(mff), 0)
  mff <- mff[c(1, m)]      
  if ("terms" %in% class(formula)){
    mff$formula <- formula
  } else {
    environment(frame.form) <- environment(formula)
    mff$formula <- frame.form
  }
  mff$drop.unused.levels <- drop.unused.levels
  mff$xlev <- xlev
  mff[[1]] <- as.name("model.frame")      
  mf <- eval(mff, parent.frame(2))  
  
  # 2. get response and wts
  # index
  index <- eval(substitute(index(da), list(da = mc$data)))
  if (match("subset", names(mff), 0)){
    idx <- eval(substitute(with(da, subset), list(da = mc$data, subset = mc$subset)))
    index <- index[idx, ]
  }
  index <- droplevels(index)
  attr(mf, "index") <- index  
  alt <- index$alt
  chid <- index$chid  
  id <- index$id
  id <- as.integer(if (!is.null(id)) id else chid) - 1
  nC <- nlevels(alt)
  # response
  y <- as.double(model.response(mf, "any"))
  # y can be null in prediction
  if (length(y)){
    attr(y, "nzero") <- as.integer(which(y == 1) - 1) #zero indexed 
    attr(y, "choice") <- as.integer(alt[y == 1])
  }
  #if (length(choice) != nlevels(chid))
  #  stop("# choices not equal to # chids.")  
  # wts
  wts <- model.weights(mf)
  if (is.null(wts)) 
    wts <- rep(1, nrow(index))  
  
  # 3. get design matrix 
  im <- as(alt, "sparseMatrix")    
  if (!isTRUE(validObject(im, test = TRUE))) 
    stop("invalid conditioning factor: ")   
  grp <- 0
  X <- Matrix(0, nrow(mf), 0)
  
  ## a) intercepts and alternative specific vars with contanst coef 
  has.int <- attr(terms(bars[[1]]), "intercept")
  if (!has.int)
    bars[[1]] <- as.formula(paste0(deparse(bars[[1]]), "+1"))
  X2 <- model.Matrix(eval(bars[[1]]), mf, sparse = TRUE, xlev = xlev,
                     drop.unused.levels = drop.unused.levels)     
  if (ncol(X2) > 0){  
    cs <- sqrt(colSums(X2^2) - nrow(X2) * colMeans(X2)^2)
    idx <- which(cs == 0)
    assign <- X2@assign
    assign <- assign[-idx]
    X2 <- X2[, -idx, drop = FALSE]
    if (has.int){
      X3 <- im
      X3@x <- rep(1, nrow(X2))
      X3 <- if (use.contrast) t(X3)[, -1] else t(X3)
      X <- cBind(X, X3)      
      grp <- c(grp, ncol(X3))    
    }
  }
  if (ncol(X2) > 0){
    assign <- getAssign(bars[[1]])[assign]
    grp2 <- cumsum(as.integer(table(assign)))        
    grp <- c(grp, grp[length(grp)] + grp2)
    X <- cBind(X, X2)
  }
  
  ## b) individual specific vars that need to vary by alternative
  if (!is.null(bars[[2]])){
    mm <- model.matrix(eval(bars[[2]]), mf)
    if (colnames(mm)[1] == "(Intercept)")
      mm <- mm[, -1, drop = FALSE]
    if (ncol(mm) > 0){
      Z <- lapply(seq_len(ncol(mm)), function(j) {
        im@x <- mm[, j]
        if (use.contrast) t(im)[, -1] else t(im)
      })  
      X <- cBind(X, do.call(cBind, Z))  
      grp <- c(grp, max(grp) + 1:ncol(mm) * ifelse(use.contrast, nrow(im) - 1, nrow(im)))  
    }          
  }   
  
  ## c) alternative specific vars with alt varying coef
  if (!is.null(bars[[3]])){
    mm <- model.matrix(eval(bars[[3]]), mf)
    if (colnames(mm)[1] == "(Intercept)")
      mm <- mm[, -1, drop = FALSE]
    if (ncol(mm) > 0){
      Z <- lapply(seq_len(ncol(mm)), function(j) {
        im@x <- mm[, j]
        t(im) 
      })  
      X <- cBind(X, do.call(cBind, Z))
      grp <- c(grp, grp[length(grp)] + 1:ncol(mm) * nrow(im))
    }
  }    
  
  # get grouping info
  gid <- as.numeric(table(chid))    
  gid <- as.integer(c(0, cumsum(gid))) 
  attr(gid, "nalt") <- as.integer(diff(gid))
  gvar <- as.integer(grp)
  attr(gvar, "nvar") <- as.integer(diff(gvar))  
  
  # return results
  list(mf = mf, X = X, y = y, wts = wts, 
       alt = alt, chid = chid, id = as.integer(id), 
       gid = gid, gvar = gvar,
       formula = formula)
}


#FIXME: now only handles factors. 
# Need to figure out a fast sparse QR to deal with the general case.
#scaleX <- function(X){  
#  cs <- sqrt(colSums(X^2)/nrow(X))
#  for (i in 1:ncol(X))  
#    X@x[(X@p[i] + 1):X@p[i + 1]] <- X@x[(X@p[i] + 1):X@p[i + 1]]/cs[i]
#  return(list(X = X, R = cs))
#}


scaleX <- function(X, gvar){
  nG <- length(gvar) - 1
  nO <- nrow(X)
  XU <- vector("list", nG)
  X2 <- Matrix(0, nrow(X), 0)
  for (i in 1:nG){
    idx <- (gvar[i] + 1):gvar[i + 1]
    XX <- X[, idx, drop = FALSE]
    XtX <- crossprod(XX)/nO
    U <- chol(XtX)
    X2 <- cBind(X2, XX %*% solve(U))
    XU[[i]] <- U
  }
  dimnames(X2) <- dimnames(X)
  return(list(X = X2, XU = XU))
}